﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Unity.Mathematics;
using UnityEditor;
using UnityEditor.Animations;
using UnityEngine;
using UnityEngine.Experimental.Rendering;

namespace Mars.GpuInstance
{
    public class GPUAnimationBaker
    {
        private bool _isBakingAnimation;
        private GPUAnimationBakeData _animationBakeData;
        private ComputeShader _skinnedMeshBakeComputeShader;
        private int _animationToTextureKernelID;

        private AnimatorController _bakerAnimatorController;
        private AnimationClip _bakerAnimationClip;

        public static readonly string COMPUTE_SKINNED_MESH_BAKE_PATH = "Compute/CSSkinnedMeshBake";
        public static readonly string COMPUTE_ANIMATION_TO_TEXTURE_KERNEL = "CSAnimationToTexture";

        public class GPUAnimationBakeData
        {
            public GPUAnimatorInstancerForBake crowdPrototype;
            public int[] textureSize;

            public RenderTexture animationRenderTexture;
            //public RenderTexture animQuaternionRT;                  // 四元数记录旋转，存储时每个维度使用 byte，总共 RGBA 32bit
            //public RenderTexture animOffsetScaleRT;                 // xyz 记录位移，w 记录 X 缩放，存储时每个维度使用 half，总共 RGBA Half
            //public RenderTexture animAdditionalScaleRT;             // xyzw 记录非等比缩放，存储时每个维度使用 8bit，两两组成 16bit 表达一个 half 数据，分别表示 YZ 缩放，总共 RGBA 32bit

            public HashSet<int> nonPropScalingBoneIndexes;          // 有非等比缩放动画的骨骼索引
            public bool currentClipHasNonPropScaling;               // 当前动画是否有非等比缩放
            
            public GameObject sampleObject;
            public Animator sampleAnimator;
            public SkinnedMeshRenderer[] sampleSkinnedMeshRenderers;
            public int currentClipIndex;
            public int currentClipFrame;
            public Matrix4x4[] boneMatrices;                        // 中间数据矩阵，使用三组数据表达变换
            public Vector4[] quaternionData;
            public Vector4[] offsetScaleData;
            public Vector4[] nonPropScaleData;
            public Matrix4x4[] anchorMatrices;
            public AnimatorOverrideController animatorOverrideController;
            public bool isBakeInitialized;
            public List<Transform> boneTransforms;
            public List<Matrix4x4> bindPoses;
            public List<int> bindPosesBoneCount;
            public Dictionary<string, AnimationClip> clipDict;
            public List<string> clipNames;
            public Action cb;
        }


        #region SkinnedMesh Bake Methods
        
        public void SkinnedMeshBakeAnimations(GPUAnimatorInstancerForBake crowdPrototype, Action cb = null)
        {
            if (crowdPrototype.gameObject == null || crowdPrototype.gameObject.GetComponentInChildren<SkinnedMeshRenderer>() == null)
            {
                return;
            }
            _animationBakeData = new GPUAnimationBakeData();
            _animationBakeData.crowdPrototype = crowdPrototype;
            _animationBakeData.cb = cb;
            EditorApplication.update += AnimationBakerInitialize;
        }

        public void AnimationBakerInitialize()
        {
            EditorApplication.update -= AnimationBakerInitialize;

            try
            {
                if (_animationBakeData.crowdPrototype.animationData == null)
                {
                    GenerateAnimationData(_animationBakeData.crowdPrototype);
                }
                _bakerAnimatorController = Resources.Load<AnimatorController>("Editor/BakerAnimator");
                _bakerAnimationClip = Resources.Load<AnimationClip>("Editor/BakerAnimation");

                _animationBakeData.sampleObject = GetSampleGameObject(_animationBakeData.crowdPrototype);
                _animationBakeData.sampleObject.transform.localScale = Vector3.one;
                _animationBakeData.sampleObject.name += "(Baker)";
                _animationBakeData.sampleObject.hideFlags = HideFlags.DontSave;

                _animationBakeData.sampleObject.transform.position = Vector3.zero;
                _animationBakeData.sampleObject.transform.rotation = Quaternion.identity;
                _animationBakeData.sampleSkinnedMeshRenderers = _animationBakeData.sampleObject.GetComponentsInChildren<SkinnedMeshRenderer>(true);
                _animationBakeData.sampleAnimator =  _animationBakeData.sampleObject.GetComponent<Animator>();
                _animationBakeData.sampleAnimator.runtimeAnimatorController = _animationBakeData.crowdPrototype.gameObject.GetComponent<Animator>().runtimeAnimatorController;
                _animationBakeData.sampleAnimator.cullingMode = AnimatorCullingMode.AlwaysAnimate;

                _animationBakeData.sampleAnimator.applyRootMotion = true;

                _animationBakeData.isBakeInitialized = false;
                _animationBakeData.boneTransforms = new List<Transform>();
                _animationBakeData.bindPoses = new List<Matrix4x4>();
                _animationBakeData.bindPosesBoneCount = new List<int>();
                _animationBakeData.clipDict = new Dictionary<string, AnimationClip>();
                _animationBakeData.clipNames = new List<string>();

                _animationBakeData.nonPropScalingBoneIndexes = new HashSet<int>();
            }
            catch (Exception e)
            {
                ClearBakingData();
                //Debug.LogError(e);
                return;
            }

            AnimationBakerStart();
        }

        public void AnimationBakerStart()
        {
            try
            {
                if (EditorUtility.DisplayCancelableProgressBar("烘焙动画", "烘焙...", 0.1f))
                {
                    ClearBakingData();
                    return;
                }

                if (_animationBakeData.sampleAnimator.runtimeAnimatorController == null)
                {
                    throw new Exception($"{_animationBakeData.crowdPrototype.name}没有挂AnimatorController");
                }

                var animationData = _animationBakeData.crowdPrototype.animationData;

                AnimatorController animatorController = null;

                if (_animationBakeData.sampleAnimator.runtimeAnimatorController is AnimatorOverrideController)
                {
                    animatorController = (AnimatorController) ((AnimatorOverrideController) _animationBakeData.sampleAnimator.runtimeAnimatorController).runtimeAnimatorController;
                }
                else
                {
                    animatorController = (AnimatorController) _animationBakeData.sampleAnimator.runtimeAnimatorController;
                }

                for (int i = 0; i < animatorController.layers.Length; i++)
                {
                    var layer = animatorController.layers[i];
                    for (int j = 0; j < layer.stateMachine.states.Length; j++)
                    {
                        var state = layer.stateMachine.states[j];
                        if (state.state.motion && state.state.motion is AnimationClip)
                        {
                            var clip = state.state.motion as AnimationClip;
                            if (clip)
                            {
                                if (_animationBakeData.clipDict.TryAdd(state.state.name, clip))
                                {
                                    _animationBakeData.clipNames.Add(state.state.name);
                                }
                                else
                                {
                                    //Debug.LogError($"存在同名动画，{state.state.name}");
                                }
                            }
                            else
                            {
                                //Debug.LogError($"添加动画失败，{state.state.name} 不是动画");
                            }
   
                        }
                    }
                }
                
                animationData.totalFrameCount = 0;

                #region Collect Skinned Mesh Renderer Data

                if (animationData.skinnedMeshDataList == null)
                {
                    animationData.skinnedMeshDataList = new List<GPUSkinnedMeshData>();
                }
                else
                {
                    animationData.skinnedMeshDataList.Clear();
                }

                List<string> dontDestroyBoneNames = new List<string>();
                if (animationData.bones == null)
                {
                    animationData.bones = new List<GPUBone>();
                }
                else
                {
                    animationData.bones.Clear();
                }

                foreach (SkinnedMeshRenderer skinnedMeshRenderer in _animationBakeData.sampleSkinnedMeshRenderers)
                {
                    if (skinnedMeshRenderer.sharedMesh == null)
                    {
                        throw new Exception($"{skinnedMeshRenderer.gameObject.name}的网格数据是空的" + skinnedMeshRenderer.gameObject.name);
                    }

                    var smd = new GPUSkinnedMeshData
                    {
                        transformName = skinnedMeshRenderer.gameObject.name,
                        boneIndexes = new int[skinnedMeshRenderer.bones.Length],
                        hasBindPoseOffset = false,
                        bindPoseOffset = Matrix4x4.identity,
                    };

                    for (int b = 0; b < skinnedMeshRenderer.bones.Length; b++)
                    {
                        Transform boneTransform = skinnedMeshRenderer.bones[b];
                        if (boneTransform == null)
                        {
                            continue;
                        }
                        var bone = animationData.GetBoneByTransform(boneTransform.name);
                        if (bone == null)
                        {
                            bone = new GPUBone
                            {
                                boneTransformName = boneTransform.name,
                                boneIndex = animationData.bones.Count,
                            };
                            animationData.bones.Add(bone);
                            _animationBakeData.boneTransforms.Add(boneTransform);
                            _animationBakeData.bindPoses.Add(skinnedMeshRenderer.sharedMesh.bindposes[b]);
                            _animationBakeData.bindPosesBoneCount.Add(skinnedMeshRenderer.bones.Length);
                        }

                        smd.boneIndexes[b] = bone.boneIndex;
                        if (_animationBakeData.bindPoses[bone.boneIndex] != skinnedMeshRenderer.sharedMesh.bindposes[b])
                        {
                            if (_animationBakeData.bindPosesBoneCount[bone.boneIndex] < skinnedMeshRenderer.bones.Length)
                            {
                                _animationBakeData.bindPoses[bone.boneIndex] = skinnedMeshRenderer.sharedMesh.bindposes[b];
                                _animationBakeData.bindPosesBoneCount[bone.boneIndex] = skinnedMeshRenderer.bones.Length;
                            }
                        }
                    }

                    animationData.skinnedMeshDataList.Add(smd);
                }
                
                foreach (SkinnedMeshRenderer skinnedMeshRenderer in _animationBakeData.sampleSkinnedMeshRenderers)
                {
                    var smd = animationData.GetSkinnedMeshDataByName(skinnedMeshRenderer.gameObject.name);
                    if (_animationBakeData.bindPoses[smd.boneIndexes[0]] != skinnedMeshRenderer.sharedMesh.bindposes[0])
                    {
                        smd.hasBindPoseOffset = true;
                        smd.bindPoseOffset = _animationBakeData.bindPoses[smd.boneIndexes[0]].inverse * skinnedMeshRenderer.sharedMesh.bindposes[0];
                    }
                }

                int smdIndex = 0;
                foreach (SkinnedMeshRenderer skinnedMeshRenderer in _animationBakeData.sampleSkinnedMeshRenderers)
                {
                    var smd = animationData.skinnedMeshDataList[smdIndex];
                    for (int i = 0; i < smd.boneIndexes.Length; i++)
                    {
                        if (skinnedMeshRenderer.bones[i] == null)
                        {
                            continue;
                        }
                        var bone = animationData.GetBoneByTransform(skinnedMeshRenderer.bones[i].name);
                        Transform parentTransform = skinnedMeshRenderer.bones[i].parent;

                        while (parentTransform != null)
                        {
                            var parentBone = animationData.GetBoneByTransform(parentTransform.name);
                            if (parentBone != null)
                            {
                                animationData.BoneSetParent(bone, parentBone);
                                break;
                            }

                            parentTransform = parentTransform.parent;
                        }
                    }

                    smdIndex++;
                }

                animationData.totalBoneCount = animationData.bones.Count;

                if (animationData.totalBoneCount == 0)
                {
                    //Debug.LogError($"{_animationBakeData.crowdPrototype.name}没有骨骼数据");
                    ClearBakingData();
                    return;
                }

                #endregion Collect Skinned Mesh Renderer Data

                #region Collect Animation Clip Data

                var previousClipDataList = animationData.clipDataList;
                animationData.clipDataList = new List<GPUAnimationClipData>();

                foreach (var clipData in _animationBakeData.clipNames)
                {
                    var animationClip = _animationBakeData.clipDict[clipData];
                    if (animationClip == null)
                    {
                        continue;
                    }
                    int frameCount = Mathf.CeilToInt(animationClip.length * _animationBakeData.crowdPrototype.frameRate + 1);

                    int isLoopDisabled = animationClip.isLooping ? 0 : 1;

                    if (previousClipDataList != null)
                    {
                        foreach (var preClipData in previousClipDataList)
                        {
                            if (preClipData.clipName == clipData)
                            {
                                isLoopDisabled = preClipData.isLoopDisabled;
                                break;
                            }
                        }
                    }

                    var acd = new GPUAnimationClipData
                    {
                        clipName = clipData,
                        clipIndex = animationData.clipDataList.Count,
                        clipStartFrame = animationData.totalFrameCount,
                        clipFrameCount = frameCount,
                        length = (frameCount - 1.0f) / _animationBakeData.crowdPrototype.frameRate,
                        isLoopDisabled = isLoopDisabled,
                    };

                    animationData.totalFrameCount += acd.clipFrameCount;
                    animationData.clipDataList.Add(acd);
                }

                #endregion Collect Animation Clip Data
                
                _isBakingAnimation = true;
                _animationBakeData.currentClipIndex = 0;
                _animationBakeData.currentClipFrame = 0;
                _animationBakeData.boneMatrices = new Matrix4x4[animationData.totalBoneCount * animationData.totalFrameCount];
                _animationBakeData.quaternionData = new Vector4[animationData.totalBoneCount * animationData.totalFrameCount];
                _animationBakeData.offsetScaleData = new Vector4[animationData.totalBoneCount * animationData.totalFrameCount];
                _animationBakeData.nonPropScaleData = new Vector4[animationData.totalBoneCount * animationData.totalFrameCount];
                _animationBakeData.anchorMatrices = new Matrix4x4[animationData.totalBoneCount * animationData.totalFrameCount];

                _animationBakeData.sampleAnimator.runtimeAnimatorController = _bakerAnimatorController;

                EditorApplication.update += AnimationBakerUpdate;
            }
            catch (Exception e)
            {
                ClearBakingData();
                //Debug.LogError(e);
            }
        }

        public void OverrideAnimator(AnimationClip animationClip)
        {
            if (_animationBakeData.animatorOverrideController == null)
                _animationBakeData.animatorOverrideController =
                    new AnimatorOverrideController(_animationBakeData.sampleAnimator.runtimeAnimatorController);

            _animationBakeData.animatorOverrideController[_bakerAnimationClip] = animationClip;

            _animationBakeData.sampleAnimator.runtimeAnimatorController =
                _animationBakeData.animatorOverrideController;
        }
        
        Vector3 GetScaleFromMatrix(Matrix4x4 matrix) => new Vector3(matrix.GetColumn(0).magnitude, matrix.GetColumn(1).magnitude, matrix.GetColumn(2).magnitude);

        public void AnimationBakerUpdate()
        {
            if (!_isBakingAnimation)
            {
                ClearBakingData();
                return;
            }

            try
            {
                var animationData = _animationBakeData.crowdPrototype.animationData;

                if (animationData.clipDataList.Count == 0)
                {
                    throw new Exception("动画列表为空");
                }

                do
                {
                    var animationClipData = animationData.clipDataList[_animationBakeData.currentClipIndex];

                    if (EditorUtility.DisplayCancelableProgressBar("烘焙动画", "帧数: " +
                        _animationBakeData.currentClipFrame + "/" + animationClipData.clipFrameCount
                        + " Clip: " + (_animationBakeData.currentClipIndex + 1) + "/" +
                        _animationBakeData.crowdPrototype.animationData.clipDataList.Count + " (" +
                        animationClipData.clipName + ")",
                        0.1f + ((_animationBakeData.currentClipFrame + animationClipData.clipStartFrame) / (float) _animationBakeData.crowdPrototype.animationData.totalFrameCount) * 0.8f))
                    {
                        ClearBakingData();
                        return;
                    }

                    #region Initialize Sample Object

                    if (!_animationBakeData.isBakeInitialized)
                    {
                        _animationBakeData.isBakeInitialized = true;

                        var animationClip = _animationBakeData.clipDict[animationClipData.clipName];
                        OverrideAnimator(animationClip);
                        _animationBakeData.sampleAnimator.Update(0);
                        return;
                    }

                    #endregion Initialize Sample Object
                    
                    _animationBakeData.sampleObject.transform.position = Vector3.zero;
                    _animationBakeData.sampleObject.transform.rotation = Quaternion.identity;

                    var hasNonPropScaling = false;
                    for (int i = 0; i < _animationBakeData.boneTransforms.Count; i++)
                    {
                        var boneTransform = _animationBakeData.boneTransforms[i];
                        
                        int boneIndex = animationData.GetBoneIndexByTransform(boneTransform.name);
                        int index = (_animationBakeData.currentClipFrame + animationClipData.clipStartFrame) * animationData.totalBoneCount + boneIndex;
                        var boneMatrix = boneTransform.localToWorldMatrix * _animationBakeData.bindPoses[boneIndex];
                        
                        var position = Vector3.zero;
                        var rotation = Quaternion.identity;
                        var scale = Vector3.one;
                        if(boneMatrix == Matrix4x4.zero) Debug.LogError("[GPUAnimationBaker] 矩阵数据异常, matrix == Matrix4x4.zero");
                        else DecomposeMatrix(boneMatrix, out position, out rotation, out scale);
                        
                        if (!(math.distance(scale.x, scale.y)< 0.01f && math.distance(scale.y, scale.z) < 0.01f))
                        {
                            hasNonPropScaling = true;
                            //Debug.Log($"[GPUAnimationBaker] non prop scaling. scale = {scale}, bone = {boneTransform}, anim = {animationClipData.clipName}, frame = {_animationBakeData.currentClipFrame}");
                            _animationBakeData.nonPropScalingBoneIndexes.Add(i);
                        }
                        _animationBakeData.boneMatrices[index] = boneMatrix;
                        _animationBakeData.quaternionData[index] = new Vector4(rotation.x, rotation.y, rotation.z, rotation.w);
                        _animationBakeData.offsetScaleData[index] = new Vector4(position.x, position.y, position.z, scale.x);
                        _animationBakeData.nonPropScaleData[index] = new Vector4(scale.y, scale.z, 0, 0);
                        _animationBakeData.anchorMatrices[index] = boneTransform.localToWorldMatrix;
                    }
                    _animationBakeData.currentClipHasNonPropScaling |= hasNonPropScaling;

                    #region Move To Next Frame

                    _animationBakeData.currentClipFrame++;
                    if (animationClipData.clipFrameCount == _animationBakeData.currentClipFrame)
                    {
                        animationClipData.nonPropScaling = _animationBakeData.currentClipHasNonPropScaling;
                        animationClipData.nonPropScaleBoneIndexes =
                            _animationBakeData.nonPropScalingBoneIndexes.ToArray();
                        
                        animationData.clipDataList[_animationBakeData.currentClipIndex] =  animationClipData;
                        _animationBakeData.currentClipHasNonPropScaling = false;
                        _animationBakeData.nonPropScalingBoneIndexes.Clear();
                        
                        _animationBakeData.currentClipIndex++;
                        _animationBakeData.currentClipFrame = 0;
                    }

                    if (_animationBakeData.currentClipIndex == _animationBakeData.crowdPrototype.animationData.clipDataList.Count)
                    {
                        EditorApplication.update -= AnimationBakerUpdate;
                        AnimationBakerFinish();
                        return;
                    }

                    var nextAnimationClip = animationData.clipDataList[_animationBakeData.currentClipIndex];
                    if (animationClipData.clipIndex != nextAnimationClip.clipIndex)
                    {
                        var animationClip = _animationBakeData.clipDict[nextAnimationClip.clipName];
                        OverrideAnimator(animationClip);
                        _animationBakeData.sampleAnimator.Play(_animationBakeData.sampleAnimator.GetCurrentAnimatorStateInfo(0).shortNameHash, 0, 0);
                        _animationBakeData.sampleAnimator.Update(0);
                        return;
                    }

                    _animationBakeData.sampleAnimator.Update(nextAnimationClip.length / (nextAnimationClip.clipFrameCount - 1));

                    #endregion Move To Next Frame
                } while (true);
            }
            catch (Exception e)
            {
                ClearBakingData();
                //Debug.LogError(e);
            }
        }

        private static readonly int OutputAnimationRT = Shader.PropertyToID("outputRT");
        private static readonly int AnimationData = Shader.PropertyToID("animationData");
        //private static readonly int OutputQuaternionRT = Shader.PropertyToID("outputQuaternionRT");
        //private static readonly int OutputOffsetScaleRT = Shader.PropertyToID("outputOffsetScaleRT");
        //private static readonly int OutputAdditionalScaleRT = Shader.PropertyToID("outputAdditionalScaleRT");
        //private static readonly int QuaternionData = Shader.PropertyToID("quaternionData");
        //private static readonly int OffsetScaleData = Shader.PropertyToID("offsetScaleData");
        //private static readonly int AdditionalScaleData = Shader.PropertyToID("additionalScaleData");
        private static readonly int dataCount = Shader.PropertyToID("dataCount");
        private static readonly int textureSizeX = Shader.PropertyToID("textureSizeX");
        //private static readonly int nonPropDataCount = Shader.PropertyToID("nonPropDataCount");
        //private static readonly int nonPropTextureSize = Shader.PropertyToID("nonPropTextureSize");

        public void AnimationBakerFinish()
        {
            try
            {
                if (_skinnedMeshBakeComputeShader == null)
                {
                    _skinnedMeshBakeComputeShader = (ComputeShader) Resources.Load(COMPUTE_SKINNED_MESH_BAKE_PATH);
                    _animationToTextureKernelID = _skinnedMeshBakeComputeShader.FindKernel(COMPUTE_ANIMATION_TO_TEXTURE_KERNEL);
                }
                
                var animationData = _animationBakeData.crowdPrototype.animationData;
                var clipDataList = animationData.clipDataList;
                var totalDataCount = 0;
                
                var animationTexData = new float4[1024 * 1024];
                for (int i = 0; i < clipDataList.Count; i++)
                {
                    var clipData = clipDataList[i];
                    clipData.animTexDataStartIndex = totalDataCount;
                    
                    // 采用优化策略，对于等比缩放的骨骼，将矩阵拆分为 位移+缩放（xyz、scale 4个数据）和旋转（四元数 4个数据），每骨骼每帧需要 8 个数据，2 个 color
                    // 对于非等比缩放的骨骼，无法拆分矩阵分量，只能继续使用矩阵运算，每骨骼每帧需要 16 个数据，4 个 color
                    var nonPropBoneCount = clipData.nonPropScaleBoneIndexes.Length;
                    var propBoneCount = animationData.totalBoneCount - nonPropBoneCount;
                    var propBoneDataCountPerFrame = propBoneCount * 2;
                    var nonPropBoneDataCountPerFrame = nonPropBoneCount * 4;
                    
                    var propBoneDataCount = propBoneDataCountPerFrame * clipData.clipFrameCount;
                    clipData.animTexNonPropDataStartIndex = clipData.animTexDataStartIndex + propBoneDataCount;
                    
                    totalDataCount += clipData.clipFrameCount * ( propBoneDataCountPerFrame + nonPropBoneDataCountPerFrame);

                    // 按照是否有非等比缩放的骨骼，将骨骼数据分为两组，一组是等比缩放的，一组是非等比缩放的
                    var boneDataArr = new GPUAnimationClipData.BoneData[animationData.totalBoneCount];
                    int normalIndex = 0, nonPropIndex = 0;
                    for (int j = 0; j < clipData.clipFrameCount; j++)
                    {
                        for (int k = 0; k < animationData.totalBoneCount; k++)
                        {
                            if (j == 0)
                            {
                                // 只判断一次骨骼是否有非等比缩放
                                var nonPropScaling = clipData.nonPropScaleBoneIndexes.Contains(k);
                                boneDataArr[k] = new GPUAnimationClipData.BoneData()
                                {
                                    nonPropScaling = nonPropScaling,
                                    index = nonPropScaling ? nonPropIndex++ : normalIndex++
                                };
                            }

                            var boneData = boneDataArr[k];
                            var bakeDataIndex = (clipData.clipStartFrame + j) * animationData.totalBoneCount + k;
                            if (!boneData.nonPropScaling)
                            {
                                var texIndex = clipData.animTexDataStartIndex + (j * propBoneCount + boneData.index)* 2;
                                animationTexData[texIndex] = _animationBakeData.quaternionData[bakeDataIndex];
                                animationTexData[texIndex + 1] = _animationBakeData.offsetScaleData[bakeDataIndex];
                            }
                            else
                            {
                                var texIndex = clipData.animTexNonPropDataStartIndex + (j * nonPropBoneCount + boneData.index) * 4;
                                animationTexData[texIndex] = _animationBakeData.boneMatrices[bakeDataIndex].GetColumn(0);
                                animationTexData[texIndex + 1] = _animationBakeData.boneMatrices[bakeDataIndex].GetColumn(1);
                                animationTexData[texIndex + 2] = _animationBakeData.boneMatrices[bakeDataIndex].GetColumn(2);
                                animationTexData[texIndex + 3] = _animationBakeData.boneMatrices[bakeDataIndex].GetColumn(3);
                            }
                        }
                    }
                    
                    
                    // 写回数据
                    clipData.boneDataArr = boneDataArr;
                    clipData.propBoneCount = propBoneCount;
                    clipDataList[i] = clipData;
                }
                

                animationData.textureSizeX = Mathf.NextPowerOfTwo(Mathf.CeilToInt(Mathf.Sqrt(totalDataCount)));
                animationData.textureSizeY = Mathf.NextPowerOfTwo(Mathf.CeilToInt(totalDataCount / (float) animationData.textureSizeX));
                _animationBakeData.textureSize = new int[2] {animationData.textureSizeX, animationData.textureSizeY};

                if (animationData.textureSizeX > 8192 || animationData.textureSizeY > 8192)
                {
                    throw new Exception("Bone and animation frame amount is too high.");
                }
                
                
                // 统计含有非等比缩放的动画总帧数和涉及的骨骼数量
                // var nonPropScalingBoneIndexes = new List<int>(_animationBakeData.nonPropScalingBoneIndexes);
                // var nonPropScalingBoneCount = nonPropScalingBoneIndexes.Count;
                // var nonPropScalingClips = new List<int>();
                // var nonPropScalingFrameCount = 0;
                // for (int i = 0; i < clipDataList.Count; i++)
                // {
                //     var clipData = clipDataList[i];
                //     if (clipData.nonPropScaling)
                //     {
                //         nonPropScalingClips.Add(i);
                //         nonPropScalingFrameCount+= clipData.clipFrameCount;
                //     }
                // }
                // var nonPropScalingDataCount = nonPropScalingBoneCount * nonPropScalingFrameCount;
                // var nonPropScalingData = new Vector4[nonPropScalingDataCount];
                //
                // // 按顺序写入每一帧的 scaleYZ 数据
                // var curFrame = 0;
                // foreach (var clipIndex in nonPropScalingClips)
                // {
                //     var animationClipData = animationData.clipDataList[clipIndex];
                //     animationClipData.clipNonPropScaleStartFrame = curFrame;
                //     animationData.clipDataList[clipIndex] = animationClipData;
                //     for (int j = 0; j < animationClipData.clipFrameCount; j++)
                //     {
                //         for (int k = 0; k < nonPropScalingBoneCount; k++)
                //         {
                //             var globalIndex = ( j + animationClipData.clipStartFrame) * animationData.totalBoneCount + nonPropScalingBoneIndexes[k];
                //             var nonPropScale = _animationBakeData.nonPropScaleData[globalIndex];
                //             
                //             var nonPropDataIndex = curFrame * nonPropScalingBoneCount + k;
                //             nonPropScalingData[nonPropDataIndex] = nonPropScale;
                //         }
                //
                //         curFrame++;
                //     }
                // }

                
                #region Create Render Texture
                // 创建旋转+位移+X缩放数据RT
                _animationBakeData.animationRenderTexture = new RenderTexture(animationData.textureSizeX, animationData.textureSizeY, 0, RenderTextureFormat.ARGBHalf)
                {
                    filterMode = FilterMode.Point,
                    enableRandomWrite = true,
                    useMipMap = false,
                    autoGenerateMips = false
                };
                _animationBakeData.animationRenderTexture.Create();
                
                // _animationBakeData.animQuaternionRT = new RenderTexture(animationData.textureSizeX, animationData.textureSizeY, 0, RenderTextureFormat.ARGBHalf)
                // {
                //     filterMode = FilterMode.Point,
                //     enableRandomWrite = true,
                //     useMipMap = false,
                //     autoGenerateMips = false
                // };
                // _animationBakeData.animOffsetScaleRT = new RenderTexture(animationData.textureSizeX, animationData.textureSizeY, 0, RenderTextureFormat.ARGBHalf)
                // {
                //     filterMode = FilterMode.Point,
                //     enableRandomWrite = true,
                //     useMipMap = false,
                //     autoGenerateMips = false
                // };
                // _animationBakeData.animQuaternionRT.Create();
                // _animationBakeData.animOffsetScaleRT.Create();
                
                // 创建非等比缩放数据RT
                // animationData.nonPropScaleTexSizeX =  Mathf.NextPowerOfTwo(Mathf.CeilToInt(Mathf.Sqrt(nonPropScalingDataCount)));
                // animationData.nonPropScaleTexSizeY = Mathf.NextPowerOfTwo(Mathf.CeilToInt(nonPropScalingDataCount / (float) animationData.nonPropScaleTexSizeX));
                // _animationBakeData.animAdditionalScaleRT = new RenderTexture(
                //     animationData.nonPropScaleTexSizeX > 0 ? animationData.nonPropScaleTexSizeX : 1,
                //     animationData.nonPropScaleTexSizeY > 0 ? animationData.nonPropScaleTexSizeY : 1, 0,
                //     RenderTextureFormat.ARGB32)
                // {
                //     filterMode = FilterMode.Point,
                //     enableRandomWrite = true,
                //     useMipMap = false,
                //     autoGenerateMips = false
                // };
                // _animationBakeData.animAdditionalScaleRT.Create();

                #endregion Create Render Texture
                
                
                // 将帧数据写入贴图
                ComputeBuffer animationTexBuffer = new ComputeBuffer(totalDataCount, sizeof(float) * 4);
                animationTexBuffer.SetData(animationTexData, 0, 0, totalDataCount);
                //var boneMatrices = _animationBakeData.boneMatrices;
                // ComputeBuffer quaternionBuffer = new ComputeBuffer(boneMatrices.Length, 16);
                // ComputeBuffer offsetScaleBuffer = new ComputeBuffer(boneMatrices.Length, 16);
                // ComputeBuffer additionalScaleBuffer = new ComputeBuffer(nonPropScalingDataCount > 0 ? nonPropScalingDataCount : 1, 16);
                
                // quaternionBuffer.SetData(_animationBakeData.quaternionData);
                // offsetScaleBuffer.SetData(_animationBakeData.offsetScaleData);
                // additionalScaleBuffer.SetData(nonPropScalingData);

                _skinnedMeshBakeComputeShader.SetTexture(_animationToTextureKernelID, OutputAnimationRT, _animationBakeData.animationRenderTexture);
                _skinnedMeshBakeComputeShader.SetBuffer(_animationToTextureKernelID, AnimationData, animationTexBuffer);
                _skinnedMeshBakeComputeShader.SetInt(dataCount, totalDataCount);
                _skinnedMeshBakeComputeShader.SetInt(textureSizeX, animationData.textureSizeX);
                _skinnedMeshBakeComputeShader.Dispatch(_animationToTextureKernelID, Mathf.CeilToInt(totalDataCount / InstancerConstants.COMPUTE_SHADER_THREAD_COUNT), 1, 1);

                animationTexBuffer.Release();
                
                // _skinnedMeshBakeComputeShader.SetTexture(_animationToTextureKernelID, OutputQuaternionRT, _animationBakeData.animQuaternionRT);
                // _skinnedMeshBakeComputeShader.SetTexture(_animationToTextureKernelID, OutputOffsetScaleRT, _animationBakeData.animOffsetScaleRT);
                // _skinnedMeshBakeComputeShader.SetTexture(_animationToTextureKernelID, OutputAdditionalScaleRT, _animationBakeData.animAdditionalScaleRT);
                // _skinnedMeshBakeComputeShader.SetBuffer(_animationToTextureKernelID, QuaternionData, quaternionBuffer);
                // _skinnedMeshBakeComputeShader.SetBuffer(_animationToTextureKernelID, OffsetScaleData, offsetScaleBuffer);
                // _skinnedMeshBakeComputeShader.SetBuffer(_animationToTextureKernelID, AdditionalScaleData, additionalScaleBuffer);
                // _skinnedMeshBakeComputeShader.SetInt(dataCount, _animationBakeData.boneMatrices.Length);
                // _skinnedMeshBakeComputeShader.SetInts(textureSize, _animationBakeData.textureSize);
                // _skinnedMeshBakeComputeShader.SetInt(nonPropDataCount, nonPropScalingDataCount);
                // _skinnedMeshBakeComputeShader.SetInts(nonPropTextureSize,
                //     new int[] { animationData.nonPropScaleTexSizeX, animationData.nonPropScaleTexSizeY });
                // _skinnedMeshBakeComputeShader.Dispatch(_animationToTextureKernelID, Mathf.CeilToInt(_animationBakeData.boneMatrices.Length / InstancerConstants.COMPUTE_SHADER_THREAD_COUNT), 1, 1);
                
                // quaternionBuffer.Release();
                // offsetScaleBuffer.Release();
                // additionalScaleBuffer.Release();
                
                if (EditorUtility.DisplayCancelableProgressBar("烘焙动画", "写入贴图...", 0.91f))
                {
                    ClearBakingData();
                    return;
                }
                
                EditorUtility.SetDirty(_animationBakeData.crowdPrototype.animationData);
                EditorUtility.SetDirty(_animationBakeData.crowdPrototype);
                AssetDatabase.SaveAssets();
                AssetDatabase.Refresh();

                animationData.animTex = SaveDataTexture(_animationBakeData.animationRenderTexture, animationData.animTex, "animMap", GraphicsFormat.R16G16B16A16_SFloat, _animationBakeData.crowdPrototype.gameObject);
                
                // animationData.animQuaternionTex = SaveDataTexture(_animationBakeData.animQuaternionRT, 
                //     animationData.animQuaternionTex, "animQuaternionMap", GraphicsFormat.R16G16B16A16_SFloat, _animationBakeData.crowdPrototype.gameObject);
                // animationData.animOffsetScaleTex = SaveDataTexture(_animationBakeData.animOffsetScaleRT, 
                //     animationData.animOffsetScaleTex, "animOffsetScaleMap", GraphicsFormat.R16G16B16A16_SFloat, _animationBakeData.crowdPrototype.gameObject);
                // if (nonPropScalingDataCount > 0)
                // {
                //     animationData.animAdditionalScaleTex = SaveDataTexture(_animationBakeData.animAdditionalScaleRT, 
                //         animationData.animAdditionalScaleTex, "animScaleYZMap",GraphicsFormat.R8G8B8A8_UNorm, _animationBakeData.crowdPrototype.gameObject);
                // }
                // else
                // {
                //     if (animationData.animAdditionalScaleTex != null)
                //     {
                //         AssetDatabase.DeleteAsset(AssetDatabase.GetAssetPath(animationData.animAdditionalScaleTex));
                //     }
                //     animationData.animAdditionalScaleTex = null;
                // }
                //
                // animationData.nonPropScaleBoneIndexes = nonPropScalingBoneIndexes;

                _animationBakeData.crowdPrototype.animationData = animationData;
                EditorUtility.SetDirty(_animationBakeData.crowdPrototype.animationData);
                EditorUtility.SetDirty(_animationBakeData.crowdPrototype);
                AssetDatabase.SaveAssets();
                AssetDatabase.Refresh();

                GenerateAnimationBoneData(_animationBakeData);

                AssetDatabase.SaveAssets();
                AssetDatabase.Refresh();

                EditorUtility.ClearProgressBar();

            }
            catch (Exception e)
            {
                ClearBakingData();
                
                Debug.LogError($"[GPUAnimationBaker] {e}");
            }
            
            ClearBakingData();
        }
        
        public void DecomposeMatrix(Matrix4x4 matrix, out Vector3 position, out Quaternion rotation, out Vector3 scale)
        {
            position = matrix.GetColumn(3);
            
            scale.x = matrix.GetColumn(0).magnitude;
            scale.y = matrix.GetColumn(1).magnitude;
            scale.z = matrix.GetColumn(2).magnitude;

            // If any scale is zero, ignore rotation part
            if (scale.x == 0 || scale.y == 0 || scale.z == 0)
            {
                rotation = Quaternion.identity;
                return;
            }

            // Create a new matrix with normalized scale
            Matrix4x4 normalizedMatrix = matrix;
            normalizedMatrix[0, 0] /= scale.x;
            normalizedMatrix[0, 1] /= scale.x;
            normalizedMatrix[0, 2] /= scale.x;

            normalizedMatrix[1, 0] /= scale.y;
            normalizedMatrix[1, 1] /= scale.y;
            normalizedMatrix[1, 2] /= scale.y;

            normalizedMatrix[2, 0] /= scale.z;
            normalizedMatrix[2, 1] /= scale.z;
            normalizedMatrix[2, 2] /= scale.z;

            rotation = normalizedMatrix.rotation;
        }

        public void ClearBakingData()
        {
            EditorApplication.update -= AnimationBakerUpdate;
            EditorUtility.ClearProgressBar();
            if (_animationBakeData != null)
            {
                if (_animationBakeData.sampleObject)
                {
                    GameObject.DestroyImmediate(_animationBakeData.sampleObject);
                }
                var cb = _animationBakeData.cb;
                if (cb != null)
                {
                    cb.Invoke();
                }
            }


            _isBakingAnimation = false;
            _animationBakeData = null;
        }

        public GameObject GetSampleGameObject(GPUAnimatorInstancerForBake crowdPrototype)
        {
            return GameObject.Instantiate(crowdPrototype.gameObject);
        }

        #endregion SkinnedMesh Bake Methods
        
        public static GPUAnimationData GenerateAnimationData(GPUAnimatorInstancerForBake prototype)
        {
            if (prototype.animationData == null)
            {
                prototype.animationData = ScriptableObject.CreateInstance<GPUAnimationData>();
            }

            if (!Application.isPlaying && string.IsNullOrEmpty(AssetDatabase.GetAssetPath(prototype.animationData)))
            {
                var assetName = InstancerEditorUtil.GetPrefabAssetName(prototype.gameObject);
                var folder = GetAssetFolder(prototype.gameObject);
                string assetPath = folder + assetName + "_animationData.asset";
                if (!System.IO.Directory.Exists(folder))
                {
                    System.IO.Directory.CreateDirectory(folder);
                }

                AssetDatabase.CreateAsset(prototype.animationData, assetPath);
                
                prototype.animationData = AssetDatabase.LoadAssetAtPath<GPUAnimationData>(assetPath);
            }
            EditorUtility.SetDirty(prototype.animationData);
            EditorUtility.SetDirty(prototype);
            return prototype.animationData;
        }

        public static void GenerateAnimationBoneData(GPUAnimationBakeData bakeData)
        {
            if (bakeData.crowdPrototype == null)
            {
                return;
            }

            var prototype = bakeData.crowdPrototype;
            if (!Application.isPlaying)
            {
                var assetName = InstancerEditorUtil.GetPrefabAssetName(prototype.gameObject);
                var folder = GetAssetFolder(prototype.gameObject);
                string assetPath = folder + assetName + "_anchor.bytes";
                if (!System.IO.Directory.Exists(folder))
                {
                    System.IO.Directory.CreateDirectory(folder);
                }

                GenerateAnimationBoneData(assetPath, bakeData.anchorMatrices);
                prototype.anchorInfo = assetPath;
                //prototype.anchorInfo = AssetDatabase.LoadAssetAtPath<TextAsset>(assetPath);
            }
            EditorUtility.SetDirty(prototype);
        }
        
        public static void GenerateAnimationBoneData(string path, Matrix4x4[] datas)
        {
            var anchorInfos = new AnchorInfo[datas.Length];
            for (int i = 0; i < datas.Length; i++)
            {
                var matrix = datas[i];
                var info = new AnchorInfo();
                info.position = matrix.GetPosition();
                info.rotation = matrix.rotation;
                anchorInfos[i] = info;
            }
            using var fs = File.OpenWrite(path);
            using var pinnedData = new QuickArray<AnchorInfo>(anchorInfos);
            using var ms = pinnedData.GetUnmanagedMemoryStream();
            ms.CopyTo(fs);
        }

        //获取资源目录
        public static string GetAssetFolder(GameObject gameObject)
        {
            var path = AssetDatabase.GetAssetPath(gameObject);
            var folder = path.Substring(0, path.LastIndexOf("/"));
            folder = folder.Substring(0, folder.LastIndexOf("/")) + "/gpuskin/";
            return folder;
        }

        public static Texture2D SaveDataTexture(RenderTexture rt, Texture2D tex, string suffix, GraphicsFormat format, GameObject prefab)
        {
            string texturePath = null;
            if (tex != null)
            {
                texturePath = AssetDatabase.GetAssetPath(tex);
            }

            tex = new Texture2D(rt.width, rt.height, format, 0, TextureCreationFlags.IgnoreMipmapLimit)
            {
                filterMode = FilterMode.Point,
            };
            Rect readingRect = new Rect(0, 0, rt.width, rt.height);
            RenderTexture.active = rt;
            tex.ReadPixels(readingRect, 0, 0);
            tex.Apply(true, true);
            RenderTexture.active = null;
            rt.Release();
                
            if (string.IsNullOrEmpty(texturePath))
            {
                var assetName = InstancerEditorUtil.GetPrefabAssetName(prefab);
                var folder = GetAssetFolder(prefab);
                texturePath = $"{folder}{assetName}_{suffix}.asset";
                    
                if (!System.IO.Directory.Exists(folder))
                {
                    System.IO.Directory.CreateDirectory(folder);
                }
            }
            
            AssetDatabase.CreateAsset(tex, texturePath);
            AssetDatabase.SaveAssets();
            AssetDatabase.Refresh();

            return AssetDatabase.LoadAssetAtPath<Texture2D>(texturePath);
        }
    }
}
